mod auth;

use crate::app::AppState;
use crate::oidc;
use crate::oidc::{AuthFlowCache, UserSessionData, AUTHORIZATION_CODE, CLIENT_ID, CLIENT_SECRET, CODE, CODE_CHALLENGE, CODE_CHALLENGE_METHOD, CODE_VERIFIER, GRANT_TYPE, NONCE, OIDC_AUTH_FLOW_STATE_KEY, OPENID, REDIRECT_URI, RESPONSE_TYPE, SCOPE, STATE, USER_SESSION_KEY};
use crate::support::error::{AppError, AppResult};
use crate::support::response::AppResponse;
use anyhow::Context;
use axum::extract::{Query, State};
use axum::response::{Html, Redirect};
use axum::Router;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tower_sessions::Session;
use url::Url;

pub fn create_router() -> Router<AppState> {
    Router::new()
        .route("/", axum::routing::get(index_handler))
        .route("/login", axum::routing::get(login_handler))
        .route("/login/oauth2/code", axum::routing::get(auth_code_callback_handler))
        .route("/auth-info", axum::routing::get(auth_info_handler))
        .fallback(async || -> AppResult<()>{
            tracing::warn!("Not found");
            Err::<(), AppError>(AppError::NotFound)
        })
        .method_not_allowed_fallback(async || -> AppResult<()>{
            tracing::warn!("Method not allowed");
            Err::<(), AppError>(AppError::MethodNotAllowed)
        })
}


// #[debug_handler]
async fn index_handler(session: Session) -> AppResult<Html<String>> {
    //1. 从Session中提取UserDataSession
    let user_session_opt: Option<UserSessionData> = session.get(USER_SESSION_KEY)
        .await?;

    if let Some(user_session) = user_session_opt {
        if user_session.claims.is_some() {
            let claims_display = format!("{:#?}", user_session.claims);
            let access_token_display = user_session.access_token.as_deref().unwrap_or("Not found");

            let content = format!(
                r#"<h1>Welcome!</h1>
                   <p>You are logged in.</p>
                   <p>ID Token Claims: <pre>{}</pre></p>
                   <p>Access Token: <pre>{}</pre></p>
                   <p><a href="/userinfo">View User Info (from UserInfo endpoint)</a></p>
                   <p><a href="/login">Login Again</a></p>
                   <p><a href="/logout">Logout</a></p>"#,
                claims_display,
                access_token_display
            );
            Ok(Html(content))
        } else {
            eprintln!("User has active session but no claims. Redirecting to login.");
            Err(AppError::AuthenticationRequired)
        }
    } else {
        println!("No active user session data found. Redirecting to login.");
        Err(AppError::AuthenticationRequired)
    }
}

// #[debug_handler]
async fn login_handler(
    session: Session,
    State(AppState { oidc_client }): State<AppState>,
) -> AppResult<Redirect> {
    //1. 是否已经登录了
    if let Some(user_data) = session.get::<UserSessionData>(USER_SESSION_KEY).await? {
        if user_data.claims.is_some() {
            return Ok(Redirect::temporary("/"));
        }
    }
    //2.没有登录，尝试登录：
    session.remove::<AuthFlowCache>(OIDC_AUTH_FLOW_STATE_KEY).await?;
    //2.1  生成pkce_challenge & nonce & csrf_state
    use oidc::pkce::*;
    let pkce_verifier = generate_pkce_verifier();
    let pkce_challenge = calculate_pkce_challenge(&pkce_verifier);
    let csrf_state = generate_random_str(32);
    let nonce = generate_random_str(32);

    //2. 放入 session 确保下次可以用作校验
    let oidc_temp_state_data = AuthFlowCache {
        state: csrf_state.clone(),
        nonce: nonce.clone(),
        pkce_verifier,
    };
    session.insert(OIDC_AUTH_FLOW_STATE_KEY, oidc_temp_state_data).await?;

    //3. 构建Oidc认证发起并重定向 (开启认证流程)
    let config = crate::config::get();
    let auth_config = config.auth_config();
    let client_id = auth_config.client_id();
    let mut auth_url = Url::parse(oidc_client.oidc_metadata.authorization_endpoint())?;
    let redirect_uri = Url::parse(auth_config.redirect_uri())?;
    auth_url.query_pairs_mut()
        .append_pair(RESPONSE_TYPE, CODE)
        .append_pair(SCOPE, OPENID)
        .append_pair(CLIENT_ID, client_id)
        .append_pair(STATE, &csrf_state)
        .append_pair(NONCE, &nonce)
        .append_pair(CODE_CHALLENGE, &pkce_challenge)
        .append_pair(REDIRECT_URI, redirect_uri.as_str())
        .append_pair(CODE_CHALLENGE_METHOD, "S256");
    let auth_redirect_url = auth_url.as_str();
    tracing::info!("Auth URL: {}", auth_url);
    Ok(Redirect::temporary(auth_redirect_url))
}

#[derive(Debug, Deserialize)]
pub struct OAuthCodeCallbackQueryParams {
    pub code: String,
    pub state: String,
}

// --- TokenResponse from Keycloak ---
#[derive(Debug, Deserialize, Serialize)]
pub struct TokenResponse {
    access_token: String,
    expires_in: u64,
    refresh_expires_in: Option<u64>,
    refresh_token: Option<String>,
    token_type: String,
    id_token: Option<String>,
    #[serde(rename = "not-before-policy")]
    not_before_policy: Option<u64>,
    scope: Option<String>,
}

// #[debug_handler]
async fn auth_code_callback_handler(
    session: Session,
    State(AppState { oidc_client }): State<AppState>,
    Query(query): Query<OAuthCodeCallbackQueryParams>,
) -> Result<Redirect, AppError> {
    tracing::info!("Received callback with query: {:?}", query);
    let auth_flow_cache: AuthFlowCache = session
        .get::<AuthFlowCache>(OIDC_AUTH_FLOW_STATE_KEY)
        .await?.ok_or_else(|| AppError::AuthFlowError("OIDC auth flow state not found in session.".to_string()))?;

    // 验证state
    if query.state != auth_flow_cache.state {
        return Err(AppError::AuthFlowError("CSRF state mismatch.".into()));
    }

    // 请求 accessToken
    let auth_config = crate::config::get().auth_config();
    let redirect_uri = auth_config.redirect_uri();
    let client_id = auth_config.client_id();
    let client_secret = auth_config.client_secret();
    let oidc_metadata = oidc_client.oidc_metadata;
    let token_endpoint = oidc_metadata.token_endpoint();
    let http_client = oidc_client.http_client;
    let mut params = HashMap::new();
    params.insert(GRANT_TYPE, AUTHORIZATION_CODE);
    params.insert(CODE, &query.code);
    params.insert(REDIRECT_URI, redirect_uri);
    params.insert(CLIENT_ID, client_id);
    params.insert(CLIENT_SECRET, client_secret);
    params.insert(CODE_VERIFIER, &auth_flow_cache.pkce_verifier);

    tracing::info!("params: {:?}",params);
    let token_response = http_client.post(token_endpoint)
        .form(&params)
        .send()
        .await
        .context("Failed to send token request")?;
    if !token_response.status().is_success() {
        let error_body = token_response.text().await.unwrap_or_else(|_| String::from("Failed to read error body from token endpoint"));
        return Err(AppError::AuthFlowError(error_body));
    }
    // 验证token
    let token_json_value: TokenResponse = token_response.json()
        .await
        .context("Failed to deserialize response body from token endpoint")?;
    let token = token_json_value.access_token;
    let claims = oidc_client.jwk_service.validate_token(&token).await?;

    // 校验id_token的nonce
    let id_token = token_json_value.id_token.ok_or(AppError::AuthenticationRequired)?;

    oidc_client.jwk_service.validate_id_token(&auth_flow_cache.nonce, &id_token).await?;

    // 登录成功，放置UserDataSession
    let user_session_data = UserSessionData {
        claims: Some(claims),
        access_token: Some(token),
        refresh_token: token_json_value.refresh_token,
    };
    session.insert(USER_SESSION_KEY, user_session_data).await?;
    session.remove::<AuthFlowCache>(OIDC_AUTH_FLOW_STATE_KEY).await?;
    tracing::info!("User successfully authenticated and session established.");
    let base_url = auth_config.base_url();
    Ok(Redirect::temporary(base_url))
}

// #[debug_handler]
async fn auth_info_handler(session: Session) -> AppResult<AppResponse<UserSessionData>> {
    if let Some(user_data) = session.get::<UserSessionData>(USER_SESSION_KEY).await? {
        if user_data.claims.is_some() {
            return Ok(AppResponse::ok("Ok", Some(user_data)));
        }
    }
    Err(AppError::AuthenticationRequired)
}